package org.nimi317.web_gis.service.impl;

import cn.dev33.satoken.stp.StpUtil;
import cn.hutool.core.util.ObjectUtil;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.baomidou.mybatisplus.core.metadata.IPage;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import org.nimi317.web_gis.Enum.ModelStatus;
import org.nimi317.web_gis.dao.ModelDao;
import org.nimi317.web_gis.entity.Lstm;
import org.nimi317.web_gis.entity.Model;
import org.nimi317.web_gis.entity.TestInfo;
import org.nimi317.web_gis.form.post.ModelPagePost;
import org.nimi317.web_gis.form.post.ModelPost;
import org.nimi317.web_gis.service.LstmService;
import org.nimi317.web_gis.service.ModelService;
import org.nimi317.web_gis.service.TestinfoService;
import org.nimi317.web_gis.utils.ModelUtils;
import org.springframework.beans.BeanUtils;
import org.springframework.context.annotation.Lazy;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.web.multipart.MultipartFile;

import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;

/**
 * 模型(Model)表服务实现类
 *
 * @author makejava
 * @since 2024-05-13 15:29:25
 */
@Service("modelService")
@Transactional
public class ModelServiceImpl extends ServiceImpl<ModelDao, Model> implements ModelService {

    private final ModelUtils modelUtils;

    private final LstmService lstmService;

    private final TestinfoService testinfoService;

    public ModelServiceImpl(@Lazy ModelUtils modelUtils, LstmService lstmService, TestinfoService testinfoService) {
        this.modelUtils = modelUtils;
        this.lstmService = lstmService;
        this.testinfoService = testinfoService;
    }

    @Override
    public void train(ModelPost post, MultipartFile file) {
        // 第一步 保存csv文件
        // 第二步 进行模型训练 训练只能存在一个线程进行训练 如果已经存在则直接保存等待时候训练
        String url = modelUtils.saveCsv(file);
        Model model = new Model();
        BeanUtils.copyProperties(post, model);
        model.setUserId((Integer) StpUtil.getSession().getLoginId());
        this.save(model);
        // 保存lstm模型参数
        AtomicInteger index = new AtomicInteger();
        Integer modelId = model.getId();
        List<Lstm> lstm = post.getLstm().stream().map(item -> item.toLstm(modelId, index.getAndIncrement())).toList();
        lstmService.saveBatch(lstm);
        if (modelUtils.train(modelId, url, post)) {
            model.setStatus(ModelStatus.Waiting);
            this.updateById(model);
        } else {
            model.setStatus(ModelStatus.Full);
            this.updateById(model);
        }
    }

    @Override
    public void handleError(Integer id) {
        Model entity = new Model();
        entity.setId(id);
        entity.setStatus(ModelStatus.Error);
        this.updateById(entity);
    }

    @Override
    public void handleSuccess(Integer id, String s) {
        Model entity = new Model();
        entity.setId(id);
        entity.setStatus(ModelStatus.Finished);
        entity.setUrl(s);
        this.updateById(entity);
    }

    @Override
    public void handleRunning(Integer id) {
        Model entity = new Model();
        entity.setId(id);
        entity.setStatus(ModelStatus.Training);
        this.updateById(entity);
    }

    @Override
    public IPage<Model> pageModel(ModelPagePost post) {
        LambdaQueryWrapper<Model> wrapper = new LambdaQueryWrapper<>();
        wrapper.eq(Model::getUserId, StpUtil.getLoginIdAsInt());
        Optional.ofNullable(post.getId())
                .ifPresent(item -> {
                    wrapper.eq(Model::getId, item);
                });
        IPage<Model> page = this.page(null, wrapper);

        // 获取模型的lstm层级信息
        List<Model> records = page.getRecords();
        if (ObjectUtil.isEmpty(records)) {
            return page;
        }

        LambdaQueryWrapper<Lstm> queryWrapper = new LambdaQueryWrapper<>();
        Map<Integer, List<Lstm>> collect = lstmService.list(queryWrapper
                .in(Lstm::getModelId, records.stream().map(Model::getId).collect(Collectors.toList()))).stream().collect(Collectors.groupingBy(Lstm::getModelId));
        page.setRecords(records.stream().peek(item -> {
            item.setLstmList(collect.get(item.getId()));
        }).toList());

        //获取模型测试集结果
        List<Integer> list = records.stream().filter(item -> item.getStatus().equals(ModelStatus.Finished)).map(Model::getId).toList();
        if (ObjectUtil.isNotEmpty(list)) {
            LambdaQueryWrapper<TestInfo> lambdaQueryWrapper = new LambdaQueryWrapper<>();
            lambdaQueryWrapper.in(TestInfo::getModelId, list);
            List<TestInfo> testInfos = testinfoService.list(lambdaQueryWrapper);
            Map<Integer, List<TestInfo>> map = testInfos.stream().collect(Collectors.groupingBy(TestInfo::getModelId));
            page.setRecords(records.stream().peek(item -> {
                if (item.getStatus().equals(ModelStatus.Finished)) {
                    item.setTestInfoList(map.get(item.getId()));
                }
            }).toList());
        }
        return page;
    }
}

